examples/speculative_decoding/example.ipynb (8)

7-8: Call out prerequisites (HF auth + git‑lfs) up front.

Daring‑Anteater via git and Llama‑3.2‑1B are likely gated; readers may need HF tokens and git‑lfs. Add a short “Prereqs” note here to reduce setup failures.

---

43-47: HF gated model note.

meta-llama/Llama-3.2-1B typically requires license acceptance and an HF token. Consider documenting HF auth/env var usage before this cell to prevent download errors.

---

62-63: Capture/verify convert() return semantics.

Some modelopt APIs return a new module; others mutate in place. Assign defensively or confirm it’s in‑place to avoid silently training the wrong module.

Apply:

-mtsp.convert(model, [("eagle", config)])
+_ret = mtsp.convert(model, [("eagle", config)])
+if _ret is not None:
+    model = _ret

---

108-113: Set a seed for reproducibility.

Add a deterministic seed to make the demo repeatable.

Apply before creating Trainer:

+transformers.set_seed(42)

---

318-320: Make cleanup idempotent.

Apply:

-%%sh
-docker rm -f trtllm_serve_spec
+%%sh
+docker rm -f trtllm_serve_spec || true

---

375-387: Propagate HF cache/token into SGLang container too (mirror TRT‑LLM).

You already mount the cache—optionally plumb the token as an env var for gated models.

Apply:

-    "-v", f"{hf_cache_dir}:/root/.cache/huggingface",
+    "-v", f"{hf_cache_dir}:/root/.cache/huggingface",
+    "-e", "HUGGINGFACE_HUB_CACHE=/root/.cache/huggingface",
+    # "-e", f"HF_TOKEN={os.environ.get('HF_TOKEN','')}",

---

459-461: Idempotent cleanup.

Apply:

-%%sh
-docker rm -f sglang_serve_spec
+%%sh
+docker rm -f sglang_serve_spec || true

---

480-480: Kernel name nit.

“modelopt+serve” is a clearer display_name since this notebook demos TRT‑LLM and SGLang, not vLLM (yet).

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

examples/speculative_decoding/example.ipynb (10)

17-19: Replace git clone with huggingface_hub to avoid git‑lfs pitfalls.

Use snapshot_download for reliability and caching.

-%%sh
-git clone https://huggingface.co/datasets/nvidia/Daring-Anteater /tmp/Daring-Anteater
+from huggingface_hub import snapshot_download
+dataset_dir = "/tmp/Daring-Anteater"
+snapshot_download(
+    repo_id="nvidia/Daring-Anteater",
+    repo_type="dataset",
+    local_dir=dataset_dir,
+    local_dir_use_symlinks=False,
+)

---

66-73: Align tokenizer limits to model; set padding side; avoid ChatML fallback.

Hardcoding 1024 may truncate; ChatML fallback can mismatch Llama‑3.2.

-tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, model_max_length=1024)
-tokenizer.pad_token_id = tokenizer.eos_token_id
-if tokenizer.chat_template is None:
-    tokenizer.chat_template = (
-        "{%- for message in messages %}"
-        "{{- '<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n' }}"
-        "{%- endfor %}"
-    )
+model_max_len = getattr(getattr(model, "config", None), "max_position_embeddings", None) or 2048
+tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, model_max_length=model_max_len)
+tokenizer.padding_side = "right"
+tokenizer.pad_token_id = tokenizer.eos_token_id
+# Avoid overriding tokenizer.chat_template with a generic template.

---

96-99: Stream dataset or use datasets library; avoid loading entire JSONL into memory.

-with open("/tmp/Daring-Anteater/train.jsonl") as f:
-    data_json = [json.loads(line) for line in f]
-train_dataset = LazySupervisedDataset(data_json[: int(len(data_json) * 0.95)], tokenizer=tokenizer)
-eval_dataset = LazySupervisedDataset(data_json[int(len(data_json) * 0.95) :], tokenizer=tokenizer)
+from datasets import load_dataset
+ds = load_dataset("json", data_files={"/all": "/tmp/Daring-Anteater/train.jsonl"})["/all"]
+splits = ds.train_test_split(test_size=0.05, seed=42)
+train_dataset = LazySupervisedDataset(splits["train"], tokenizer=tokenizer)
+eval_dataset = LazySupervisedDataset(splits["test"], tokenizer=tokenizer)

---

102-113: Don’t subclass TrainingArguments for minor toggles; pass flags directly.

-from dataclasses import dataclass, field
-@dataclass
-class TrainingArguments(transformers.TrainingArguments):
-    dataloader_drop_last: bool = field(default=True)
-    bf16: bool = field(default=True)
-
-
-training_args = TrainingArguments(
+TrainingArguments = transformers.TrainingArguments
+training_args = TrainingArguments(
     output_dir="/tmp/eagle_bf16",
     num_train_epochs=4,
     per_device_train_batch_size=1,
     per_device_eval_batch_size=1,
+    bf16=True,
+    dataloader_drop_last=True,
 )

---

122-122: Avoid private API Trainer._move_model_to_device().

Let Trainer handle placement or call model.to(...) explicitly.

-trainer._move_model_to_device(model, trainer.args.device)
+# Let Trainer manage device placement.

---

173-181: Bound max tokens/seq to tokenizer/model; don’t hardcode 8192.

-    --max_num_tokens 8192 \
-    --max_seq_len 8192 \
+    --max_num_tokens {tokenizer.model_max_length} \
+    --max_seq_len {tokenizer.model_max_length} \

---

220-247: Harden TRT‑LLM docker run: shm, HF cache/token, cleanup.

-# Generate a unique container name so we can stop/remove it later
-container_name = "trtllm_serve_spec"
+container_name = "trtllm_serve_spec"
+subprocess.call(["docker","rm","-f",container_name])
@@
-    "--shm-size=2g",
+    "--shm-size=32g",
@@
-    "-v",
-    "/tmp:/tmp",
+    "-v", "/tmp:/tmp",
+    "-v", f"{os.path.expanduser('~')}/.cache/huggingface:/root/.cache/huggingface",
+    "-e", "HUGGINGFACE_HUB_CACHE=/root/.cache/huggingface",
+    "-e", f"HF_TOKEN={os.environ.get('HF_TOKEN','')}",

What is the latest stable NGC tag for nvcr.io/nvidia/tensorrt-llm/release and does trtllm-serve accept --backend pytorch with speculative_config?

---

282-303: Remove non‑standard chat_template and add robust HTTP error handling.

-    "temperature": 0,
-    "chat_template": tokenizer.chat_template,
+    "temperature": 0,
@@
-response = requests.post(
-    "http://localhost:8000/v1/chat/completions", headers=headers, data=json.dumps(payload)
-)
-output = response.json()
-
-print(output)
+resp = requests.post(
+    "http://localhost:8000/v1/chat/completions",
+    headers=headers,
+    json=payload,
+    timeout=60,
+)
+resp.raise_for_status()
+print(resp.json())

---

423-444: Add error handling to SGLang request.

-response = requests.post(
-    "http://localhost:30000/v1/chat/completions", headers=headers, data=json.dumps(payload)
-)
-output = response.json()
-
-print(output)
+resp = requests.post(
+    "http://localhost:30000/v1/chat/completions",
+    headers=headers,
+    json=payload,
+    timeout=60,
+)
+resp.raise_for_status()
+print(resp.json().get("choices", [{}])[0].get("message", {}).get("content"))

---

39-60: Do not mutate EAGLE3_DEFAULT_CFG in place; deep‑copy before updates.

Global mutation risks bleed‑through across runs/imports.

Note: I’m aware of the prior learning that ALGO_TO_CONFIG yields deep copies in tests; here we’re using EAGLE3_DEFAULT_CFG directly, so copy locally to be safe.

-from modelopt.torch.speculative.config import EAGLE3_DEFAULT_CFG
+from modelopt.torch.speculative.config import EAGLE3_DEFAULT_CFG
+import copy
@@
-# Read Default Config for EAGLE3
-config = EAGLE3_DEFAULT_CFG["config"]
+# Read Default Config for EAGLE3 (local copy)
+config = copy.deepcopy(EAGLE3_DEFAULT_CFG["config"])

examples/speculative_decoding/example.ipynb (2)

375-387: Pass HF token into SGLang container and pre‑remove any stale container.

-container_name = "sglang_serve_spec"
+container_name = "sglang_serve_spec"
+subprocess.call(["docker","rm","-f",container_name])
@@
-    "-v", f"{hf_cache_dir}:/root/.cache/huggingface",
+    "-v", f"{hf_cache_dir}:/root/.cache/huggingface",
+    "-e", "HUGGINGFACE_HUB_CACHE=/root/.cache/huggingface",
+    "-e", f"HF_TOKEN={os.environ.get('HF_TOKEN','')}",

---

93-94: Import OK — eagle_utils present at examples/speculative_decoding/eagle_utils.py

Defines LazySupervisedDataset and DataCollatorWithPadding, so "from eagle_utils import ..." works when example.ipynb is run from that directory. If the notebook will be executed from other CWDs or the module reused elsewhere, either add examples/speculative_decoding to PYTHONPATH, convert the folder to a package and use a relative import, or install the module.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.393Z
Learning: In the TensorRT-Model-Optimizer codebase, both EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG in config.py already use deepcopied configurations (either directly via deepcopy() or through variables that were already deepcopied), making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.393Z
Learning: In the TensorRT-Model-Optimizer codebase, the EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG configurations in config.py are already deepcopied, so additional deepcopy calls are not needed when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.393Z
Learning: In the TensorRT-Model-Optimizer codebase, EAGLE1_DEFAULT_CFG in config.py already uses deepcopy(default_eagle_config) in its definition, making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

examples/speculative_decoding/example.ipynb (9)

17-18: Avoid git‑LFS headaches: use huggingface_hub instead of git clone.

Replace shell clone with snapshot_download for reliability and caching.

-!git clone https://huggingface.co/datasets/nvidia/Daring-Anteater /tmp/Daring-Anteater
+from huggingface_hub import snapshot_download
+snapshot_download(
+    repo_id="nvidia/Daring-Anteater",
+    repo_type="dataset",
+    local_dir="/tmp/Daring-Anteater",
+    local_dir_use_symlinks=False,
+)

---

64-72: Align tokenizer limits to the model and set padding side explicitly.

-# Prepare Tokenizer
-tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, model_max_length=1024)
-tokenizer.pad_token_id = tokenizer.eos_token_id
+# Prepare Tokenizer
+cfg = transformers.AutoConfig.from_pretrained(base_model)
+max_len = int(getattr(cfg, "max_position_embeddings", 2048))
+tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, model_max_length=max_len)
+tokenizer.pad_token_id = tokenizer.eos_token_id
+tokenizer.padding_side = "right"
 if tokenizer.chat_template is None:

---

95-98: Avoid loading the entire JSONL into memory; stream or use datasets.

-with open("/tmp/Daring-Anteater/train.jsonl") as f:
-    data_json = [json.loads(line) for line in f]
-train_dataset = LazySupervisedDataset(data_json[: int(len(data_json) * 0.95)], tokenizer=tokenizer)
-eval_dataset = LazySupervisedDataset(data_json[int(len(data_json) * 0.95) :], tokenizer=tokenizer)
+def stream_jsonl(path):
+    with open(path) as f:
+        for line in f:
+            yield json.loads(line)
+all_data = list(stream_jsonl("/tmp/Daring-Anteater/train.jsonl"))
+split = int(len(all_data) * 0.95)
+train_dataset = LazySupervisedDataset(all_data[:split], tokenizer=tokenizer)
+eval_dataset = LazySupervisedDataset(all_data[split:], tokenizer=tokenizer)

(Or switch to datasets.load_dataset("json") with train_test_split.)

---

101-112: Don’t subclass TrainingArguments for minor toggles.

-from dataclasses import dataclass, field
+from dataclasses import dataclass, field  # (no longer needed if not subclassing)
@@
-@dataclass
-class TrainingArguments(transformers.TrainingArguments):
-    dataloader_drop_last: bool = field(default=True)
-    bf16: bool = field(default=True)
+TrainingArguments = transformers.TrainingArguments
@@
 training_args = TrainingArguments(
     output_dir="/tmp/eagle_bf16",
     num_train_epochs=4,
     per_device_train_batch_size=1,
     per_device_eval_batch_size=1,
+    bf16=True,
+    dataloader_drop_last=True,
 )

---

121-121: Remove private API call Trainer._move_model_to_device().

This is unstable across HF releases; let Trainer handle device placement.

-trainer._move_model_to_device(model, trainer.args.device)
+# Let Trainer manage device placement

---

172-180: Bind TRT‑LLM token/sequence limits to the model.

 trtllm_serve_script = f"""trtllm-serve {base_model} \
@@
-    --max_num_tokens 8192 \
-    --max_seq_len 8192 \
+    --max_num_tokens {tokenizer.model_max_length} \
+    --max_seq_len {tokenizer.model_max_length} \

---

219-246: Container ergonomics: mount HF cache/token, bump shm, and pre‑remove old container.

 import subprocess
 import threading
+import os
@@
-# Generate a unique container name so we can stop/remove it later
+# Generate a unique container name so we can stop/remove it later
 container_name = "trtllm_serve_spec"
+home_dir = os.path.expanduser("~")
+hf_cache_dir = os.path.join(home_dir, ".cache", "huggingface")
+os.makedirs(hf_cache_dir, exist_ok=True)
+subprocess.call(["docker", "rm", "-f", container_name])

 docker_cmd = [
     "docker",
     "run",
     "--rm",
     "--net",
     "host",
-    "--shm-size=2g",
+    "--shm-size=32g",
     "--ulimit",
     "memlock=-1",
     "--ulimit",
     "stack=67108864",
+    "--ipc=host",
     "--gpus",
     "all",
-    "-v",
-    "/tmp:/tmp",
+    "-v", "/tmp:/tmp",
+    "-v", f"{hf_cache_dir}:/root/.cache/huggingface",
+    "-e", "HUGGINGFACE_HUB_CACHE=/root/.cache/huggingface",
+    "-e", "HF_HUB_ENABLE_HF_TRANSFER=1",
+    # optionally: "-e", f"HF_TOKEN={os.environ.get('HF_TOKEN','')}",
     "--name",
     container_name,
     "nvcr.io/nvidia/tensorrt-llm/release:1.1.0rc2",
     "bash",
     "-c",
     "bash /tmp/trtllm_serve.sh",
 ]

---

281-301: Harden the TRT‑LLM request; remove non‑standard fields.

 payload = {
     "model": base_model,
@@
-    "temperature": 0,
-    "chat_template": tokenizer.chat_template,
+    "temperature": 0,
 }
 headers = {"Content-Type": "application/json", "Accept": "application/json"}
 
-response = requests.post(
-    "http://localhost:8000/v1/chat/completions", headers=headers, data=json.dumps(payload)
-)
-output = response.json()
-
-print(output)
+resp = requests.post(
+    "http://localhost:8000/v1/chat/completions",
+    headers=headers,
+    json=payload,
+    timeout=60,
+)
+resp.raise_for_status()
+print(resp.json())

---

421-442: Add basic error handling on the SGLang request.

-#Send request to the SGLang server
-response = requests.post(
-    "http://localhost:30000/v1/chat/completions", headers=headers, data=json.dumps(payload)
-)
-output = response.json()
-
-print(output)
+# Send request to the SGLang server
+resp = requests.post(
+    "http://localhost:30000/v1/chat/completions",
+    headers=headers,
+    json=payload,
+    timeout=60,
+)
+resp.raise_for_status()
+print(resp.json())

examples/speculative_decoding/example.ipynb (5)

34-39: Don’t mutate imported defaults; work on a copy of EAGLE3 config.

Mutating EAGLE3_DEFAULT_CFG in place can leak across cells/runs. Use a local deepcopy.

Using your retrieved learning: tests may not need extra deepcopy, but in a long‑lived notebook kernel avoiding global mutation is safer.

 import transformers
+import copy
 
 import modelopt.torch.opt as mto
 import modelopt.torch.speculative as mtsp
 from modelopt.torch.speculative.config import EAGLE3_DEFAULT_CFG
@@
-# Read Default Config for EAGLE3
-config = EAGLE3_DEFAULT_CFG["config"]
+# Read Default Config for EAGLE3 (local copy to avoid global mutation)
+config = copy.deepcopy(EAGLE3_DEFAULT_CFG["config"])

Also applies to: 48-59

---

43-46: Note: base model may require an HF token.

Llama 3.2 repos often need auth. Add a short note or check for HF_TOKEN/env login to reduce friction.

 # Load original HF model
 base_model = "meta-llama/Llama-3.2-1B"
+import os
+if not os.environ.get("HF_TOKEN"):
+    print("Hint: this model may require Hugging Face auth (HF_TOKEN or `huggingface-cli login`).")
 model = transformers.AutoModelForCausalLM.from_pretrained(
     base_model, torch_dtype="auto", device_map="cuda"
 )

---

362-405: SGLang container: ensure HF cache env, optional token, and cleanup.

 import subprocess
 import threading
 import os
@@
 container_name = "sglang_serve_spec"
@@
-docker_cmd = [
+subprocess.call(["docker", "rm", "-f", container_name])
+docker_cmd = [
     "docker", "run",
     "--rm",
     "--net", "host",
     "--shm-size=32g",
     "--gpus", "all",
     "-v", f"{hf_cache_dir}:/root/.cache/huggingface",
     "-v", "/tmp:/tmp",
     "--ipc=host",
+    "-e", "HUGGINGFACE_HUB_CACHE=/root/.cache/huggingface",
+    "-e", "HF_HUB_ENABLE_HF_TRANSFER=1",
+    # optionally: "-e", f"HF_TOKEN={os.environ.get('HF_TOKEN','')}",
     "--name", container_name,
-    "lmsysorg/sglang:latest",
+    "lmsysorg/sglang:latest",
     "bash", "-c", "bash /tmp/sglang_serve.sh",
 ]

Also consider pinning the image tag instead of “latest” for reproducibility.

---

271-273: Optional: add readiness polling instead of manual log watching.

Example snippet to insert before sending requests:

import time, requests
def wait_ready(url, timeout=120):
    t0=time.time()
    while time.time()-t0<timeout:
        try:
            r=requests.get(url, timeout=2)
            if r.status_code in (200,404): return True
        except Exception: pass
        time.sleep(2)
    raise TimeoutError(f"Service not ready: {url}")

wait_ready("http://localhost:8000")   # TRT‑LLM
# wait_ready("http://localhost:30000") # SGLang

Also applies to: 411-413

---

113-120: Verify only the EAGLE head is trainable.

Sanity‑check that base model params are frozen post‑convert; prevents unintended fine‑tuning.

# Diagnostic check (optional cell)
trainable = [n for n,p in model.named_parameters() if p.requires_grad]
print(f"Trainable params: {len(trainable)}")
assert all("eagle" in n.lower() for n in trainable), "Unexpected trainable params outside EAGLE head."

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.393Z
Learning: In the TensorRT-Model-Optimizer codebase, both EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG in config.py already use deepcopied configurations (either directly via deepcopy() or through variables that were already deepcopied), making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.393Z
Learning: In the TensorRT-Model-Optimizer codebase, the EAGLE1_DEFAULT_CFG and EAGLE3_DEFAULT_CFG configurations in config.py are already deepcopied, so additional deepcopy calls are not needed when using ALGO_TO_CONFIG mapping in tests.

Learnt from: yeyu-nvidia
PR: NVIDIA/TensorRT-Model-Optimizer#295
File: tests/gpu/torch/speculative/plugins/test_speculative_megatron_modules.py:35-39
Timestamp: 2025-09-05T19:10:36.393Z
Learning: In the TensorRT-Model-Optimizer codebase, EAGLE1_DEFAULT_CFG in config.py already uses deepcopy(default_eagle_config) in its definition, making additional deepcopy calls unnecessary when using ALGO_TO_CONFIG mapping in tests.